import numpy as np
from enum import Enum

class SensorType(Enum):
    GPS = 1
    IMU = 2
    ODOM = 3

class KalmanFilter:
    def __init__(self):
        # 状态向量 [x, y, theta, vx, vy, vtheta]
        self.x = np.zeros((6, 1))
        
        # 状态协方差矩阵
        self.P = np.eye(6) * 1000.0
        
        # 状态转移矩阵 (离散时间)
        self.F = np.eye(6)
        
        # 过程噪声协方差矩阵
        self.Q = np.eye(6) * 0.1
        
        # 测量矩阵 - 根据不同传感器动态设置
        self.H = None
        
        # 测量噪声协方差矩阵 - 根据不同传感器动态设置
        self.R = None
        
        # 最后更新时间
        self.last_time = None
        
        # 传感器权重
        self.gps_weight = 0.0
        self.imu_weight = 0.0
        self.odom_weight = 0.0
    
    def update_time(self, current_time):
        """更新时间间隔并调整状态转移矩阵"""
        if self.last_time is None:
            self.last_time = current_time
            return
            
        dt = current_time - self.last_time
        self.last_time = current_time
        
        # 更新状态转移矩阵 (假设匀速运动模型)
        self.F[0, 3] = dt  # x += vx * dt
        self.F[1, 4] = dt  # y += vy * dt
        self.F[2, 5] = dt  # theta += vtheta * dt
        
        # 更新过程噪声矩阵
        self.Q = np.diag([
            0.1 * dt,     # x 位置噪声
            0.1 * dt,     # y 位置噪声
            0.5 * dt,     # theta 角度噪声
            0.2 * dt,     # vx 速度噪声
            0.2 * dt,     # vy 速度噪声
            0.5 * dt      # vtheta 角速度噪声
        ])
    
    def predict(self):
        """执行卡尔曼滤波预测步骤"""
        self.x = self.F @ self.x
        self.P = self.F @ self.P @ self.F.T + self.Q
        return self.x.copy(), self.P.copy()
    
    def update(self, z, sensor_type, covariance):
        """执行卡尔曼滤波更新步骤"""
        # 根据传感器类型设置测量矩阵和权重
        if sensor_type == SensorType.GPS:
            # GPS测量位置 (x, y)
            self.H = np.zeros((2, 6))
            self.H[0, 0] = 1  # 测量x
            self.H[1, 1] = 1  # 测量y
            
            # 根据GPS精度设置权重
            accuracy = covariance[0, 0]
            self.gps_weight = max(0.0, 1.0 - accuracy / 100.0)
            
            # 设置测量噪声
            self.R = np.diag([
                0.01 / (self.gps_weight + 0.001),  # x测量噪声
                0.01 / (self.gps_weight + 0.001)   # y测量噪声
            ])
            
        elif sensor_type == SensorType.IMU:
            # IMU测量角度和角速度 (theta, vtheta)
            self.H = np.zeros((2, 6))
            self.H[0, 2] = 1  # 测量theta
            self.H[1, 5] = 1  # 测量vtheta
            
            # 根据IMU精度设置权重
            accuracy = covariance[0, 0]
            self.imu_weight = max(0.0, 1.0 - accuracy / 100.0)
            
            # 设置测量噪声
            self.R = np.diag([
                0.05 / (self.imu_weight + 0.001),   # theta测量噪声
                0.01 / (self.imu_weight + 0.001)    # vtheta测量噪声
            ])
            
        elif sensor_type == SensorType.ODOM:
            # 里程计测量位置和角度 (x, y, theta)
            self.H = np.zeros((3, 6))
            self.H[0, 0] = 1  # 测量x
            self.H[1, 1] = 1  # 测量y
            self.H[2, 2] = 1  # 测量theta
            
            # 根据里程计精度设置权重
            accuracy = covariance[0, 0]
            if accuracy < 0.1:  # 机器人静止
                self.odom_weight = 0.0
            else:
                self.odom_weight = min(1.0, accuracy / 10.0)  # 缩放权重
                
            # 设置测量噪声
            self.R = np.diag([
                0.1 / (self.odom_weight + 0.001),   # x测量噪声
                0.1 / (self.odom_weight + 0.001),   # y测量噪声
                0.1 / (self.odom_weight + 0.001)    # theta测量噪声
            ])
        
        # 计算卡尔曼增益
        S = self.H @ self.P @ self.H.T + self.R
        K = self.P @ self.H.T @ np.linalg.inv(S)
        
        # 更新状态估计
        y = z - self.H @ self.x  # 测量残差
        
        # 角度归一化到[-pi, pi]
        if sensor_type in [SensorType.IMU, SensorType.ODOM]:
            while y[0, 0] > np.pi:
                y[0, 0] -= 2 * np.pi
            while y[0, 0] < -np.pi:
                y[0, 0] += 2 * np.pi
                
        self.x = self.x + K @ y
        self.P = (np.eye(self.x.shape[0]) - K @ self.H) @ self.P
        
        return self.x.copy(), self.P.copy()

class SensorFusion:
    def __init__(self):
        self.kalman_filter = KalmanFilter()
        self.last_gps_time = None
        self.last_imu_time = None
        self.last_odom_time = None
    
    def process_gps(self, pose_msg, timestamp):
        """处理GPS数据"""
        # 提取位置和协方差
        x = pose_msg.pose.pose.position.x
        y = pose_msg.pose.pose.position.y
        
        # 协方差矩阵 (6x6)
        covariance = np.array(pose_msg.pose.covariance).reshape(6, 6)
        
        # 转换为测量向量
        z = np.array([[x], [y]])
        
        # 更新滤波器时间
        self.kalman_filter.update_time(timestamp)
        
        # 预测步骤
        self.kalman_filter.predict()
        
        # 更新步骤
        state, _ = self.kalman_filter.update(z, SensorType.GPS, covariance)
        
        # 记录时间
        self.last_gps_time = timestamp
        
        return state
    
    def process_imu(self, pose_msg, timestamp):
        """处理IMU数据"""
        # 提取角度和协方差
        q = pose_msg.pose.pose.orientation
        siny_cosp = 2 * (q.w * q.z + q.x * q.y)
        cosy_cosp = 1 - 2 * (q.y * q.y + q.z * q.z)
        theta = np.arctan2(siny_cosp, cosy_cosp)
        
        # 假设角速度可以从协方差矩阵估计
        vtheta = np.sqrt(pose_msg.pose.covariance[35])  # 假设角速度方差在covariance[35]
        
        # 协方差矩阵
        covariance = np.array(pose_msg.pose.covariance).reshape(6, 6)
        
        # 转换为测量向量
        z = np.array([[theta], [vtheta]])
        
        # 更新滤波器时间
        self.kalman_filter.update_time(timestamp)
        
        # 预测步骤
        self.kalman_filter.predict()
        
        # 更新步骤
        state, _ = self.kalman_filter.update(z, SensorType.IMU, covariance)
        
        # 记录时间
        self.last_imu_time = timestamp
        
        return state
    
    def process_odom(self, pose_msg, timestamp):
        """处理里程计数据"""
        # 提取位置、角度和协方差
        x = pose_msg.pose.pose.position.x
        y = pose_msg.pose.pose.position.y
        
        q = pose_msg.pose.pose.orientation
        siny_cosp = 2 * (q.w * q.z + q.x * q.y)
        cosy_cosp = 1 - 2 * (q.y * q.y + q.z * q.z)
        theta = np.arctan2(siny_cosp, cosy_cosp)
        
        # 协方差矩阵
        covariance = np.array(pose_msg.pose.covariance).reshape(6, 6)
        
        # 转换为测量向量
        z = np.array([[x], [y], [theta]])
        
        # 更新滤波器时间
        self.kalman_filter.update_time(timestamp)
        
        # 预测步骤
        self.kalman_filter.predict()
        
        # 更新步骤
        state, _ = self.kalman_filter.update(z, SensorType.ODOM, covariance)
        
        # 记录时间
        self.last_odom_time = timestamp
        
        return state
    
    def get_current_state(self):
        """获取当前状态估计"""
        return self.kalman_filter.x.copy()    